﻿#!/usr/bin/env python
# coding: utf-8
from SocketServer import BaseRequestHandler, ThreadingUDPServer
from cStringIO import StringIO
import os
import re, socket, struct, binascii, time
import sqlite3
import threading

class Common(object):
    '''global config object'''
    
    def __init__(self):
        self.host = r'127.0.0.1'
        self.port = 53
        self.dns_server = r'202.102.152.3'
        self.remote_dns_server = r'8.8.8.8'
        self.hosts_data = os.path.abspath('.') + r'\\data.db'
        self.enable_cache = True

# init config
Common = Common()
def log_infor(str,level=0):
    if level in (1, 2, 3, 4, 5):
        print 'level', level, ':', str
    else:
        print str

def DNS_starter():

    def convert_host2ip(host2ip):
        hostlines = []
        for ip, host in host2ip:
            Family = isIp(ip)
            if Family:
                packed_ip = inet_pton(Family, ip)
                hostRegEX = ToRegEX(host)
                hostlines.append([packed_ip, hostRegEX])
        return hostlines

    def isIp(host):
        if re.match(r'^((25[0-5]|2[0-4]\d|[0-1]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[0-1]?\d\d?)$', host):
            return socket.AF_INET
        elif re.match(r'^\s*((([0-9A-Fa-f]{1,4}:){7}(([0-9A-Fa-f]{1,4})|:))|(([0-9A-Fa-f]{1,4}:){6}(:|((25[0-5]|2[0-4]\d|[01]?\d{1,2})(\.(25[0-5]|2[0-4]\d|[01]?\d{1,2})){3})|(:[0-9A-Fa-f]{1,4})))|(([0-9A-Fa-f]{1,4}:){5}((:((25[0-5]|2[0-4]\d|[01]?\d{1,2})(\.(25[0-5]|2[0-4]\d|[01]?\d{1,2})){3})?)|((:[0-9A-Fa-f]{1,4}){1,2})))|(([0-9A-Fa-f]{1,4}:){4}(:[0-9A-Fa-f]{1,4}){0,1}((:((25[0-5]|2[0-4]\d|[01]?\d{1,2})(\.(25[0-5]|2[0-4]\d|[01]?\d{1,2})){3})?)|((:[0-9A-Fa-f]{1,4}){1,2})))|(([0-9A-Fa-f]{1,4}:){3}(:[0-9A-Fa-f]{1,4}){0,2}((:((25[0-5]|2[0-4]\d|[01]?\d{1,2})(\.(25[0-5]|2[0-4]\d|[01]?\d{1,2})){3})?)|((:[0-9A-Fa-f]{1,4}){1,2})))|(([0-9A-Fa-f]{1,4}:){2}(:[0-9A-Fa-f]{1,4}){0,3}((:((25[0-5]|2[0-4]\d|[01]?\d{1,2})(\.(25[0-5]|2[0-4]\d|[01]?\d{1,2})){3})?)|((:[0-9A-Fa-f]{1,4}){1,2})))|(([0-9A-Fa-f]{1,4}:)(:[0-9A-Fa-f]{1,4}){0,4}((:((25[0-5]|2[0-4]\d|[01]?\d{1,2})(\.(25[0-5]|2[0-4]\d|[01]?\d{1,2})){3})?)|((:[0-9A-Fa-f]{1,4}){1,2})))|(:(:[0-9A-Fa-f]{1,4}){0,5}((:((25[0-5]|2[0-4]\d|[01]?\d{1,2})(\.(25[0-5]|2[0-4]\d|[01]?\d{1,2})){3})?)|((:[0-9A-Fa-f]{1,4}){1,2})))|(((25[0-5]|2[0-4]\d|[01]?\d{1,2})(\.(25[0-5]|2[0-4]\d|[01]?\d{1,2})){3})))(%.+)?\s*$', host):
            return socket.AF_INET6
        else:
            return None

    #win32 下实现 inet_pton 和 inet_ntop
    def inet_ntop(family, ipstr):   #ipv4/ipv6 integer to address(with len in it)
        if family== socket.AF_INET:
            return socket.inet_ntoa(ipstr)
        elif family== socket.AF_INET6:
            v6addr = ':'.join(('%02X%02X' % (ord(i), ord(j)))
                             for i,j in zip(ipstr[::2], ipstr[1::2]))
            return v6addr

    def inet_pton(family, addr):    #ipv4/ipv6 address to integer(with len in it)
        if family== socket.AF_INET:
            return socket.inet_aton(addr)
        elif family== socket.AF_INET6:
            if '.' in addr: # a v4 addr
                 v4addr = addr[addr.rindex(':')+1:]
                 v4addr = socket.inet_aton(v4addr)
                 v4addr = map(lambda x: ('%02X' % ord(x)), v4addr)
                 v4addr.insert(2, ':')
                 newaddr = addr[:addr.rindex(':')+1] + ''.join(v4addr)
                 return self.inet_pton(family, newaddr)
            dbyts = [0]* 8 # 8 groups
            grps = addr.split(':')
            for i,v in enumerate(grps):
                if v:
                    dbyts[i] = int(v, 16)
                else:
                    for j, w in enumerate(grps[::-1]):
                        if w:
                            dbyts[7-j] = int(w, 16)
                        else:
                            break
                    break
            return ''.join( (chr(i//256) + chr(i%256)) for i in dbyts)
        else:
            raise

    def ToRegEX(hostname):
        return '^' + hostname.replace('.', '\.').replace('*','[A-Za-z0-9._-]*').replace('?', '[A-Za-z0-9._-]') + '$'


    class Struct(object): # this isnt socket.struct, this is just for a return list
        def __init__(self, **kwargs):
            for name, value in kwargs.items():
                setattr(self, name, value)

    def parse_dns_message(data):
        message = StringIO(data)
        message.seek(4)     # skip id, flag
        c_qd, c_an, c_ns, c_ar = struct.unpack('!4H', message.read(8))
        # parse question
        question = parse_dns_question(message)
        for i in range(1, c_qd):    # skip other question
            parse_dns_question(message)
        records = []
        for i in range(c_an+c_ns):
            record = parse_dns_record(message)
            if record:
                records.append(record)
        return Struct(question=question, records=records)

    def parse_dns_question(message):
        qname = parse_domain_name(message)
        qtype, qclass = struct.unpack('!HH', message.read(4))
        end_offset = message.tell()
        return Struct(name=qname, type_=qtype, class_=qclass, end_offset=end_offset)

    def parse_dns_record(message):
        domain_name = parse_domain_name(message)    # skip name in question and read names in answer
        AnswerType = struct.unpack('!h', message.read(2))[0] # read type. 1=ipv4 28=ipv6 5=CNAME
        message.seek(2, os.SEEK_CUR)                # skip class
        ttl_offset = message.tell()
        ttl = struct.unpack('!I', message.read(4))[0]
        rd_len = struct.unpack('!H', message.read(2))[0]
        ip = ''
        if AnswerType == 1:     # read content with ipv4,ipv6 and CNAME in it
            ip = inet_ntop(socket.AF_INET, message.read(rd_len))
        elif AnswerType == 28:
            ip = inet_ntop(socket.AF_INET6, message.read(rd_len))
            ip = ip.replace('0000', ':').replace(':::::', '::').replace('::::', '::').replace(':::','::') # zip IPV6 address
        else:
            message.seek(rd_len, os.SEEK_CUR)
            return False
        return Struct(ttl_offset=ttl_offset, ttl=ttl, domain_name=domain_name, ip=ip)

    def _parse_domain_labels(message):
        labels = []
        len = ord(message.read(1)) # 域名第一部分的长度。ord为转换字符类型为二进制数，相反的函数为chr
        while len > 0:
            if len >= 64:   # domain name compression
                len = len & 0x3f # 如果域名被压缩了就把长度截断为64一下，这个作用是获取高位
                offset = (len << 8) + ord(message.read(1)) # 将超过64的hex转换为数字
                mesg = StringIO(message.getvalue())
                mesg.seek(offset)
                labels.extend(_parse_domain_labels(mesg))
                return labels
            else:
                labels.append(message.read(len))
                len = ord(message.read(1))
        return labels

    def parse_domain_name(message):
        return '.'.join(_parse_domain_labels(message))

    def get_response(data):
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) # socket for the remote DNS server
        sock.connect((Common.dns_server, 53))
        sock.sendall(data)
        while True:
            sock.settimeout(0.2)
            try:
                rspdata = sock.recv(65535)
                rsp = parse_dns_message(rspdata)
                hostname = rsp.question.name.strip()
                for record in rsp.records:
                    if (record.ip in blackhole):
                        log_infor ("notice:[ %s >> %s ] have been GFWed!" %(hostname, record.ip))
                        rspdata = get_response_remote(data)
                        break
                    else:
                        log_infor ("infor: [ %s >> %s ] correctly!" %(hostname, record.ip))
                break
            except socket.timeout:
                log_infor ("notice:China DNS can\'t find the hostname, trying google DNS")
                rspdata = get_response_remote(data)
                break
            except socket.error as e:
                if e.errno != 11:
                    raise
        sock.close()
        return rspdata

    def get_response_remote(data):
        ''' get the DNS result from local DNS server '''
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.connect((Common.remote_dns_server, 53))
        sock.sendall(data)
        while True:
            sock.settimeout(0.4)
            try:
                rspdata = sock.recv(65535)
                rsp = parse_dns_message(rspdata)
                hostname = rsp.question.name
                for record in rsp.records:
                    if (record.ip in blackhole):
                        log_infor ("notice:[ %s >> %s ] have been GFWed!" %(hostname, record.ip))
                        rspdata = ''
                        break
                    else:
                        log_infor ("infor: [ %s >> %s ] correctly!" %(hostname, record.ip))
            except socket.timeout:
                log_infor ("error: [ %s ] response read timeout!" %(hostname))
                break
            except socket.error as err:
                log_infor ("%s ignored" %err.message)
        sock.close()
        return rspdata

    def update_ttl(reqdata, cache_entry):
        rspdata, cache_time = cache_entry.rspdata, cache_entry.cache_time
        rspbytes = bytearray(rspdata)
        rspbytes[:2] = reqdata[:2]          # update id
        current_time = int(time.time())
        time_interval = current_time - cache_time
        rsp = parse_dns_message(rspdata)
        for record in rsp.records:
            if record.ttl <= time_interval:
                return None
            rspbytes[record.ttl_offset:record.ttl_offset+4] = struct.pack('!I', record.ttl-time_interval)
        return str(rspbytes)

    class DNSProxyHandler(BaseRequestHandler):
        cache = {}
        def handle(self):
            reqdata, sock = self.request
            req = parse_dns_message(reqdata)
            q = req.question
            if q.type_ in (1, 28) and (q.class_ == 1): # DNS_TYPE_A = 1, DNS_TYPE_AAAA = 28, DNS_CLASS_IN = 1
                for packed_ip, hostRegEX in host_lines:
                    if re.match(hostRegEX, q.name):
                        # header, qd=1, an=1, ns=0, ar=0
                        rspdata = reqdata[:2] + '\x81\x80\x00\x01\x00\x01\x00\x00\x00\x00'
                        rspdata += reqdata[12:q.end_offset]
                        # answer
                        rspdata += '\xc0\x0c'   # pointer to domain name
                        # type, 1 for ip4, 28 for ip6
                        if len(packed_ip) == 4:
                            rspdata += '\x00\x01'   # 1 for ip4
                        else:
                            rspdata += '\x00\x1c'   # 28 for ip6
                        # class: 1, ttl: 2000(0x000007d0)
                        rspdata += '\x00\x01\x00\x00\x07\xd0'
                        rspdata += '\x00' + chr(len(packed_ip)) # rd_len
                        rspdata += packed_ip
                        sock.sendto(rspdata, self.client_address)
                        return

            # lookup cache
            if Common.enable_cache:
                cache_key = (q.name, q.type_, q.class_)
                cache_entry = self.cache.get(cache_key)
                if cache_entry:
                    rspdata = update_ttl(reqdata, cache_entry)
                    if rspdata:
                        sock.sendto(rspdata, self.client_address)
                        return
            rspdata = get_response(reqdata)
            
            if Common.enable_cache:
                self.cache[cache_key] = Struct(rspdata=rspdata, cache_time=int(time.time()))
            sock.sendto(rspdata, self.client_address)


    def blackhole_detect():
        log_infor ("Initializing BLACKHOLE list")
        # init BLACKHOLE list
        for x in range(64):
            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            if ((x % 2) == 1):
                sock.connect((Common.dns_server, 53))
            else:
                sock.connect((Common.remote_dns_server, 53))
            sock.sendall('\x00\xba\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x03www\x07twitter\x03com\x00\x00\x01\x00\x01')
            sock.settimeout(10)
            try:
                testholerps = sock.recv(65535)
                testhole = parse_dns_message(testholerps)
                for record in testhole.records:
                    ip = record.ip
                    if ip.split(".")[0] == "199": 
                        continue
                    if (ip and not (ip in blackhole)):
                        blackhole.append(ip)
                        log_infor ('added into blackhole list [%s]' %ip)
            except socket.error as err:
                log_infor ("failed!")
                if err.errno != 11:
                    raise
            sock.close()
        log_infor ("BLACKHOLE list initialized. bye!")

    # Initializing DNS server
    blackhole = []
    host_ip_db = sqlite3.connect(Common.hosts_data)
    host2ip = host_ip_db.cursor().execute(r'select ip,hostname from [host2ip]').fetchall()
    host_lines = convert_host2ip(host2ip)
    
    threading.Thread(target=blackhole_detect).start()
    threading.Thread(target=ThreadingUDPServer((Common.host, Common.port), DNSProxyHandler).serve_forever()).start()

def main():
    threading.Thread(target=DNS_starter).start()

if __name__ == '__main__':
    main()
